import numpy as np
import torch
import torch.nn as nn

from common_utils.utils import l2_normalize

cuda = True if torch.cuda.is_available() else False
class MLP_Generator(nn.Module):
    def __init__(self,opt):
        super(MLP_Generator, self).__init__()
        # 为类别标签创建嵌入层
        self.label_emb = nn.Embedding(opt.n_classes, opt.attSize)
        # 第一层线性层
        self.l1 = nn.Sequential(nn.Linear(opt.attSize, 2048))


    def forward(self, noise, labels):
        # 将标签嵌入到噪声中
        gen_input = torch.mul(self.label_emb(labels), noise)
        # 通过第一层线性层
        out = self.l1(gen_input)
        # 重新整形为合适的形状
        # out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        return out


class MLP_Discriminator(nn.Module):
    def __init__(self,opt):
        super(MLP_Discriminator, self).__init__()
        # Output layers
        # 第一个输出 (name=generation) 是鉴别是否认为所显示的图像是伪造的，
        self.adv_layer = nn.Sequential(nn.Linear(2048, 1), nn.Sigmoid())
        # 而第二个输出 (name=auxiliary) 是鉴别认为图像所属的类。
        self.aux_layer = nn.Sequential(nn.Linear(2048, opt.n_classes), nn.Softmax())

    def forward(self, img):
        out = l2_normalize(img)
        validity = self.adv_layer(out)
        label = self.aux_layer(out)

        return validity, label